# -*- coding: utf-8 -*-
# @Time    : 2021/9/30 10:29
# @Author  : Jonny
# @File    : LuckDbMysqlHelperUtils.py

import pymysql
from timeit import default_timer
import logging


# ---- 使用 with 的方式来优化代码
class LuckDbMysqlHelperUtils(object):

    def __init__(self, db_host, db_port, db_name, db_user, db_password, commit=True, log_time=True, log_label='总用时'):
        self.db_host = db_host
        self.db_port = db_port
        self.db_name = db_name
        self.db_user = db_user
        self.db_password = db_password
        """
        :param commit: 是否在最后提交事务(设置为False的时候方便单元测试)
        :param log_time:  是否打印程序运行总时间
        :param log_label:  自定义log的文字
        """
        self._log_time = log_time
        self._commit = commit
        self._log_label = log_label

    # 创建数据库连接
    def get_connection(self):
        conn = pymysql.connect(host=self.db_host, port=self.db_port, db=self.db_name,
                               user=self.db_user, password=self.db_password)
        return conn

    def __enter__(self):

        # 如果需要记录时间
        if self._log_time is True:
            self._start = default_timer()

        # 在进入的时候自动获取连接和cursor
        conn = self.get_connection()
        # 创建光标对象
        cursor = conn.cursor(pymysql.cursors.DictCursor)
        conn.autocommit = False

        self._conn = conn
        self._cursor = cursor
        return self

    def __exit__(self, *exc_info):
        # 提交事务
        if self._commit:
            self._conn.commit()
        # 在退出的时候自动关闭连接和cursor
        self._cursor.close()
        self._conn.close()

        if self._log_time is True:
            diff = default_timer() - self._start
            logging.info(f'-- {self._log_label}: {diff} 秒')

    @property
    def cursor(self):
        return self._cursor
